// ===== 类型定义 =====

pub enum Ast {
  Atom(String)
  List(Array[Ast])
} derive(Eq, Show)

pub enum Value {
  Symbol(String)
  Number(Int)
  Pair(Value, Value)
  Nil
  Bool(Bool)
  Lambda(Array[String], Ast, Env)
  Label(String, Value)
} derive(Eq)
typealias Array[(String, Value)] as Env

// 用于模拟 OCaml ref
struct Ref[T] {
  mut val : T
}

// ===== 词法分析器 =====

enum Token {
  LParen
  RParen
  AtomToken(String)
  QuoteToken
} derive(Show)

fn is_space(c : Char) -> Bool {
  c == ' ' || c == '\t' || c == '\n' || c == '\r'
}

fn[T] drop_while(arr : Array[T], p : (T) -> Bool) -> Array[T] {
  match arr {
    [] => []
    [x, ..xs] => if p(x) { drop_while(xs.to_array(), p) } else { arr }
  }
}

fn[T] break_on(arr : Array[T], p : (T) -> Bool) -> (Array[T], Array[T]) {
  fn loop_ (acc : Array[T], rest : Array[T]) -> (Array[T], Array[T]) {
    match rest {
      [] => (acc.rev(), [])
      [x, ..xs] => {
        if !p(x) {
          loop_([x, ..acc], xs.to_array())
        } else {
          (acc.rev(), rest)
        }
      }
    }
  }
  loop_([], arr)
}

fn tokenize_(chars : Array[Char]) -> Array[Token] {
  match chars {
    [] => []
    ['(', ..cs] => [LParen, ..tokenize_(drop_while(cs.to_array(), is_space))]
    [')', ..cs] => [RParen, ..tokenize_(drop_while(cs.to_array(), is_space))]
    ['\'', ..cs] => [QuoteToken, ..tokenize_(drop_while(cs.to_array(), is_space))]
    cs => {
      let (atom, rest) = break_on(cs, fn(c) { is_space(c) || c == '(' || c == ')' || c == '\'' })
      [AtomToken(@string.from_array(atom)), ..tokenize_(drop_while(rest, is_space))]
    }
  }
}
pub fn tokenize(s : String) -> Array[Token] {
  let t = tokenize_(s.to_array())
  t
}

// ===== 语法分析器 =====

fn parse_list(tokens : Array[Token], acc : Array[Ast]) -> (Array[Ast], Array[Token]) raise {
  let mut rest = tokens
  let nodes = acc
  while !rest.is_empty() {
    match rest[0] {
      RParen => {
        return (nodes, rest[1:].to_array())
      }
      _ => {
        let (ast, new_rest) = parse_expr(rest)
        nodes.push(ast)
        rest = new_rest
      }
    }
  }
  fail("")
}

fn parse_expr(tokens : Array[Token]) -> (Ast, Array[Token]) raise {
  if tokens.is_empty() { fail("")  }
  match tokens[0] {
    LParen => {
      let (list, rest) = parse_list(tokens[1:].to_array(), [])
      (List(list), rest)
    }
    RParen => fail("")
    AtomToken(s) => (Atom(s), tokens[1:].to_array())
    QuoteToken => {
      let (quoted, rest) = parse_expr(tokens[1:].to_array())
      (List([Atom("quote"), quoted]), rest)
    }
  }
}

pub fn parse(tokens : Array[Token]) -> Ast raise {
  let (ast, rest) = parse_expr(tokens)
  if !rest.is_empty() {
    fail("")
  }
  ast
}

// ===== 辅助函数 =====

fn try_int(s : String) -> Value {
  match (try? @strconv.parse_int(s)) {
    Ok(n) => Number(n)
    Err(_) => Symbol(s)
  }
}

fn is_true(v : Value) -> Bool {
  match v {
    Bool(false) | Nil => false
    _ => true
  }
}

fn ast_to_value(ast : Ast) -> Value {
  match ast {
    Atom(s) => try_int(s)
    List(lst) => {
      fn build_list(arr : Array[Ast]) -> Value {
        if arr.is_empty() { Nil }
        else { Pair(ast_to_value(arr[0]), build_list(arr[1:].to_array())) }
      }
      build_list(lst)
    }
  }
}

pub fn initial_env() -> Env {
  []
}

// ===== 求值器 =====

pub fn eval(ast : Ast, env : Ref[Env]) -> Value raise {
  match ast {
    Atom(s) => {
      for v in env.val {
        let (name, val) = v
        if name == s { return val }
      }
      try_int(s)
    }
    List([]) => Nil
    List([Atom("quote"), arg]) => ast_to_value(arg)
    List([Atom("atom"), arg]) => {
      let v = eval(arg, env)
      match v {
        Symbol(_) | Number(_) | Nil => Bool(true)
        _ => Bool(false)
      }
    }
    List([Atom("eq"), lhs, rhs]) => {
      fn eq(lhs : Value, rhs : Value) -> Value {
        match (lhs, rhs) {
          (Symbol(a), Symbol(b)) => Bool(a == b)
          (Number(a), Number(b)) => Bool(a == b)
          (Nil, Nil) => Bool(true)
          (Pair(a1, b1), Pair(a2, b2)) => {
            if eq(a1, a2) == Bool(true) { eq(b1, b2) } else { Bool(false) }
          }
          _ => Bool(false)
        }
      }
      eq(eval(lhs, env), eval(rhs, env))
    }
    List([Atom("car"), arg]) => {
      match eval(arg, env) {
        Pair(car, _) => car
        Nil => Nil
        _ => fail("")
      }
    }
    List([Atom("cdr"), arg]) => {
      match eval(arg, env) {
        Pair(_, cdr) => cdr
        Nil => Nil
        _ => fail("")
      }
    }
    List([Atom("cons"), lhs, rhs]) => {
      Pair(eval(lhs, env), eval(rhs, env))
    }
    List([Atom("cond"), ..clauses]) => {
      fn eval_clauses(cs : Array[Ast]) -> Value raise {
        if cs.is_empty() { fail("") }
          match cs[0] {
            List([test_, result]) => {
              let test_val = eval(test_, env)
              if is_true(test_val) { eval(result, env) } else { eval_clauses(cs[1:].to_array()) }
            }
            _ => fail("")
        }
      }
      eval_clauses(clauses.to_array())
    }
    List([Atom("lambda"), arg, body]) => {
      match arg {
        List(params) => {
          let param_names : Array[String] = []
          for p in params {
            match p {
              Atom(s) => param_names.push(s)
              _ => fail("")
            }
          }
          Lambda(param_names, body, env.val)
        }
        _ => fail("")
      }
    }
    List([Atom("label"), Atom(name), lambda]) => {
      Label(name, eval(lambda, env))
    }
    List([Atom("define"), Atom(name), expr]) => {
      let value = eval(expr, env)
      env.val = [(name, value), ..env.val]
      value
    }
    List([Atom("+"), lhs, rhs]) => {
      match (eval(lhs, env), eval(rhs, env)) {
        (Number(a), Number(b)) => Number(a + b)
        _ => fail("")
      }
    }
    List([Atom("-"), lhs, rhs]) => {
      match (eval(lhs, env), eval(rhs, env)) {
        (Number(a), Number(b)) => Number(a - b)
        _ => fail("")
      }
    }
    List([Atom("*"), lhs, rhs]) => {
      match (eval(lhs, env), eval(rhs, env)) {
        (Number(a), Number(b)) => Number(a * b)
        _ => fail("")
      }
    }
    List([Atom("/"), lhs, rhs]) => {
      match (eval(lhs, env), eval(rhs, env)) {
        (Number(a), Number(b)) => Number(a / b)
        _ => fail("")
      }
    }
    List([func, ..args]) => {
      let func_val = eval(func, env)
      let arg_vals = args.map(fn(arg) { eval(arg, env) })
      match func_val {
        Lambda(params, body, closure) => {
          let new_env: Ref[Env] = {val: closure}
          for i in 0..<params.length() {
            new_env.val = [(params[i], arg_vals[i]), ..new_env.val]
          }
          eval(body, new_env)
        }
        Label(name, Lambda(params, body, closure)) => {
          let new_env: Ref[Env] = {val: [(name, Label(name, Lambda(params, body, closure))), ..closure]}
          for i in 0..<params.length() {
            new_env.val = [(params[i], arg_vals[i]), ..new_env.val]
          }
          eval(body, new_env)
        }
        _ => fail("")
      }
    }
  }
}
// ===== 打印函数 =====

pub fn string_of_value(v : Value) -> String {
  match v {
    Symbol(s) => s
    Number(n) => n.to_string()
    Bool(true) => "t"
    Bool(false) => "()"
    Nil => "()"
    Pair(car, cdr) => {
      fn print_list(v : Value) -> String {
        match v {
          Pair(hd, tl) => " " + string_of_value(hd) + print_list(tl)
          Nil => ""
          other => " . " + string_of_value(other)
        }
      }
      "(" + string_of_value(car) + print_list(cdr) + ")"
    }
    Lambda(params, _, _) => "lambda " + params.join(",")
    Label(name, _) => "label " + name
  }
}
// ===== 组合函数 =====

pub fn s_expr_to_value(s : String, env : Ref[Env]) -> Value raise {
  let tokens = tokenize(s)
  let ast = parse(tokens)
  eval(ast, env)
}

// ===== 测试函数 =====

pub fn test_builtins() -> Unit {
  let env: Ref[Env] = {val: initial_env()}

  let tests : Array[(String, String)] = [
    ("(quote a)", "a"),
    ("(quote (a b c))", "(a b c)"),
    ("(atom 'a)", "t"),
    ("(atom '(a b))", "()"),
    ("(eq 'a 'a)", "t"),
    ("(eq 'a 'b)", "()"),
    ("(car '(a b c))", "a"),
    ("(cdr '(a b c))", "(b c)"),
    ("(cons 'a '(b c))", "(a b c)"),
    ("(cond ((eq 1 2) 'a) ((eq 1 1) 'b))", "b"),
    ("()", "()"),
    ("'(1 2 3)", "(1 2 3)"),
    ("'(a b c)", "(a b c)"),
    ("(cons 'a (cons 'b ()))", "(a b)"),
    ("(cond (() 'a) ('t 'b))", "b"),
    ("(+ 3 5)", "8"),
    ("((lambda (x) (cons 'a x)) '(b c))", "(a b c)"),
    ("((lambda (x y) (cons x (cons y ()))) 'a 'b)", "(a b)"),
    ("((lambda (x) (cons (car x) (cdr x))) '(a b c))", "(a b c)"),
    ("((lambda (f x) (f x)) (lambda (y) (cons 'a y)) '(b c))", "(a b c)"),
    ("(define add3 (lambda (x) (+ x 3)))", "lambda x"),
    ("(add3 5)", "8"),
    ("(define prod (label prod. (lambda (n) (cond ((eq n 0) 1) ('t (* n (prod. (- n 1))))))))", "label prod."),
    ("(prod 4)", "24"),
    ("(define Z (lambda (f) ((lambda (x) (f (lambda (v) ((x x) v)))) (lambda (x) (f (lambda (v) ((x x) v)))))))", "lambda f"),
    ("(define fib. (lambda (g) (lambda (n) (cond ((eq n 0) 0) ((eq n 1) 1) ('t (+ (g (- n 1)) (g (- n 2))))))))", "lambda g"),
    ("(define fib (Z fib.))", "lambda n"),
    ("(fib 6)", "8"),
  ]

  for val in tests {
    let (expr, expected) = val
      try {
        let result = string_of_value(s_expr_to_value(expr, env))
        println("测试: \{expr} 期望: \{expected} 结果: \{result}")
        if result == expected {
          println("  [通过]")
        } else {
          println("  [失败]")
        }
      }
      catch {
        e => println(" [出错] \{e}")
      }
  }
}
fn main {
  test_builtins()
}